#! /usr/bin/env python
# -*- coding: utf-8 -*-

# Copyright (c) JoinQuant Development Team
# Author: Huayong Kuang <kuanghuayong@joinquant.com>

import multiprocessing as mp
import numpy as np
import tempfile

def funcA(X):
    print("A {}".format(X.shape))
    # do something with X
    print('funcA OK')


def funcB(X):
    print("B {}".format(X.shape))
    # do something else with X
    print('funcB OK')


if __name__=='__main__':
    dim = int(2.75e8)

    pool = mp.Pool(2)

    with tempfile.NamedTemporaryFile(dir='/tmp', delete=False) as tmpfile:
        X = np.memmap(tmpfile, shape=dim, dtype=np.float32)  # create the memory map
        # init the map by chunks of size 1e8
        max_chunk_size = int(1e8)
        for start_pos in range(0, dim, max_chunk_size):
            chunk_size = min(dim-start_pos, max_chunk_size)
            X[start_pos:start_pos+chunk_size] = np.random.rand(chunk_size,)
        P=[mp.Process(target=funcA, args=(X,)), mp.Process(target=funcB, args=(X,))]
        for p in P:
            p.start()

        for p in P:
            p.join()

        pool.map(funcA, [X, X, X])

    pool.close()
    pool.join()

